
#
# A Production Rule System implemented in Python.
#
# (C) Copyright 2006-2008, Johan Lindberg
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 59 Temple
# Place, Suite 330, Boston, MA 02111-1307 USA
#

import copy
import logging

def _gen_id():
    i = 0
    while True:
        i+= 1
        yield i
        
class Token(dict):
    def __init__(self):
        self.id = _gen_id()
        self.timestamp = 0
        
#    def __repr__(self):
#        _result = ["Token: timestamp = %s; bindings = {" % (self.timestamp)]
#        for key in self.keys():
#            _result.append(" %s = %s," % (key, repr(self[key])))
#        _result.append(" }")
#        
#        return "".join(_result)
        
class Node(object):
    def __init__(self):
        self.nexts = []
        self.logger = logging.getLogger('PyRete')
        
### 
### Alpha Nodes
### -----------
### There are two types of alpha nodes (one-input nodes), one for testing
### attribute values against constants and another to test that two (or
### more) attributes in the same object are equal.
### 

def make_ConstantTestNode(attribute, relation, constant):
    logger = logging.getLogger('PyRete')
    logger.debug("make_ConstantTestNode(%s %s %s)" % (attribute, relation, constant))
    
    if relation == "==":
        class ConstantTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("ConstantTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute, relation, constant, tag, timestamp, fact))
                
                if getattr(fact, attribute) == constant:
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == "!=":
        class ConstantTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("ConstantTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute, relation, constant, tag, timestamp, fact))
                
                if getattr(fact, attribute) != constant:
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == "<":
        class ConstantTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("ConstantTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute, relation, constant, tag, timestamp, fact))
                
                if getattr(fact, attribute) < constant:
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == "<=":
        class ConstantTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("ConstantTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute, relation, constant, tag, timestamp, fact))
                
                if getattr(fact, attribute) <= constant:
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == ">":
        class ConstantTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("ConstantTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute, relation, constant, tag, timestamp, fact))
                
                if getattr(fact, attribute) > constant:
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == ">=":
        class ConstantTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("ConstantTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute, relation, constant, tag, timestamp, fact))
                
                if getattr(fact, attribute) >= constant:
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    else:
        raise Exception("Unknown comparison: %s" % (relation))
        
        
    return ConstantTestNode()
    
    
def make_AttributeTestNode(attribute1, relation, attribute2):
    logger = logging.getLogger('PyRete')
    logger.debug("make_AttributeTestNode(%s %s %s)" % (attribute1, relation, attribute2))
    
    if relation == "!=":
        class AttributeTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("AttributeTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute1, relation, attribute2, tag, timestamp, fact))
                
                if getattr(fact, attribute1) != getattr(fact, attribute2):
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == "<":
        class AttributeTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("AttributeTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute1, relation, attribute2, tag, timestamp, fact))
                
                if getattr(fact, attribute1) < getattr(fact, attribute2):
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == "<=":
        class AttributeTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("AttributeTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute1, relation, attribute2, tag, timestamp, fact))
                
                if getattr(fact, attribute1) <= getattr(fact, attribute2):
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == ">":
        class AttributeTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("AttributeTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute1, relation, attribute2, tag, timestamp, fact))
                
                if getattr(fact, attribute1) > getattr(fact, attribute2):
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    elif relation == ">=":
        class AttributeTestNode(Node):
            def right_activate(self, tag, timestamp, fact):
                self.logger.debug("AttributeTestNode(%s %s %s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attribute1, relation, attribute2, tag, timestamp, fact))
                
                if getattr(fact, attribute1) >= getattr(fact, attribute2):
                    for next in self.nexts:
                        next.right_activate(tag, timestamp, fact)
                        
    else:
        raise Exception("Unknown comparison: %s" % (relation))
        
        
    return AttributeTestNode()
    
    
def make_BindingsTestNode(*attributes):
    logger = logging.getLogger('PyRete')
    logger.debug("make_BindingsTestNode(%s)" % (", ".join(attributes)))
    
    class BindingsTestNode(Node):
        def right_activate(self, tag, timestamp, fact):
            self.logger.debug("BindingsTestNode(%s).right_activate(tag = %s, timestamp = %s, fact = %s)" % (attributes, tag, timestamp, fact))
            value, assigned, failed = None, False, False
            for attr in attributes:
                if not assigned:
                    value = getattr(fact, attr)
                    assigned = True
                    
                else:
                    if getattr(fact, attr) != value:
                        failed = True
                        break
                        
            if not failed:
                for next in self.nexts:
                    next.right_activate(tag, timestamp, fact)
                    
    return BindingsTestNode()
    
    
### Beta Nodes
### ----------
### There are two types of Beta Nodes, one for testing attributes bound to
### the same variable (Join Node) and one for testing that no facts exist
### in Working Memory that passes the join test.

## Node for testing variable bindings across two facts. This type of node is
## described in chapter 2.5.3 (The Ordinary Two-input Node) of Forgy's
## thesis[1].

## Node for testing the existance of facts allowing variable bindings across
## two facts. This type of node is described in chapter 2.5.4 (The <NOT>
## Node) of Forgy's thesis[1].

class LeftInputAdapter(Node):
    """
    The LeftInputAdapter is used to allow an AlphaNode to propagate facts
    to the first JoinNode in the network. It creates a token for each fact
    passed to it.
    """
    
    def __init__(self, binding):
        self.logger = logging.getLogger('PyRete')
        self.binding = binding
        self.nexts = []
        self.tokens = self.tokeniter()
        
    def tokeniter(self):
        for fact in self._tokens:
            token = Token()
            token[self.binding] = fact
            
            yield token
            
    def right_activate(self, tag, timestamp, fact):
        self.logger.debug("LeftInputAdapter.right_activate(tag = %s, timestamp = %s, fact = %s)" % (tag, timestamp, fact))
        
        token = Token()
        token[self.binding] = fact
        
        for next in self.nexts:
            next.left_activate(tag, timestamp, token)
            
            
class DummyJoinNode(Node):
    """
    A DummyJoinNode is used to add facts to a token without performing any
    tests on attributes.
    """
    
    def __init__(self, binding):
        self.logger = logging.getLogger('PyRete')
        self.logger.debug("DummyJoinNode.__init__(binding = %s)" % (binding))
        
        self.binding = binding
        
        self.nexts = []
        
    def left_activate(self, tag, timestamp, token):
        self.logger.debug("DummyJoinNode.left_activate(tag = %s, timestamp = %s, token = %s)" % (tag, timestamp, token))
        
        for fact in self.right.tokens:
            t = copy.copy(token)
            t[self.binding] = fact
            
            for next in self.nexts:
                next.left_activate(tag, timestamp, t)
                
    def right_activate(self, tag, timestamp, fact):
        self.logger.debug("DummyJoinNode.right_activate(tag = %s, timestamp = %s, fact = %s)" % (tag, timestamp, fact))
        
        for token in self.left.tokens:
            t = copy.copy(token)
            t[self.binding] = fact
            
            for next in self.nexts:
                next.left_activate(tag, timestamp, t)
                
                
class JoinNode(Node):
    def __init__(self, leftFact, leftAttr, rel, rightFact, rightAttr):
        self.logger = logging.getLogger('PyRete')
        self.logger.debug("JoinNode.__init__(leftFact = %s, leftAttr = %s, rel = %s, rightFact = %s, rightAttr = %s)" % (leftFact, leftAttr, rel, rightFact, rightAttr))
        
        self.testfunc = {   "==" :  lambda x, y: x == y,
                            "!=" :  lambda x, y: x != y,
                            ">=" :  lambda x, y: x >= y,
                            ">"  :  lambda x, y: x > y,
                            "<=" :  lambda x, y: x <= y,
                            "<"  :  lambda x, y: x < y, }
        self.tests = [(leftFact, leftAttr, rel, rightFact, rightAttr),]
        self.nexts = []
        
    def add_test(self, leftFact, leftAttr, rel, rightFact, rightAttr):
        self.logger.debug("JoinNode.add_test(leftFact = %s, leftAttr = %s, rel = %s, rightFact = %s, rightAttr = %s)" % (leftFact, leftAttr, rel, rightFact, rightAttr))
        
        self.tests.append((leftFact, leftAttr, rel, rightFact, rightAttr))
        
    def left_activate(self, tag, timestamp, token):
        self.logger.debug("JoinNode.left_activate(tag = %s, timestamp = %s, token = %s)" % (tag, timestamp, token))
        
        for fact in self.right.tokens:
            failed = False
            for leftFact, leftAttr, rel, rightFact, rightAttr in self.tests:
                leftValue = getattr(token[leftFact], leftAttr)
                rightValue = getattr(fact, rightAttr)
                
                if not self.testfunc[rel](leftValue, rightValue):
                    failed = True
                    break
                    
            if not failed:
                t = copy.copy(token)
                t[rightFact] = fact
                
                for next in self.nexts:
                    next.left_activate(tag, timestamp, t)
                    
    def right_activate(self, tag, timestamp, fact):
        self.logger.debug("JoinNode.right_activate(tag = %s, timestamp = %s, fact = %s)" % (tag, timestamp, fact))
        
        for token in self.left.tokens:
            failed = False
            for leftFact, leftAttr, rel, rightFact, rightAttr in self.tests:
                leftValue = getattr(token[leftFact], leftAttr)
                rightValue = getattr(fact, rightAttr)
                
                if not self.testfunc[rel](leftValue, rightValue):
                    failed = True
                    break
                    
            if not failed:
                t = copy.copy(token)
                t[rightFact] = fact
                
                for next in self.nexts:
                    next.left_activate(tag, timestamp, t)
                    
                    
class NotNode(Node):
    def __init__(self):
        self.logger = logging.getLogger('PyRete')
        self.logger.debug("NotNode.__init__()")
        
        self.testfunc = {   "==" :  lambda x, y: x == y,
                            "!=" :  lambda x, y: x != y,
                            ">=" :  lambda x, y: x >= y,
                            ">"  :  lambda x, y: x > y,
                            "<=" :  lambda x, y: x <= y,
                            "<"  :  lambda x, y: x < y, }
        self.count = {}
        
        self.tests = []
        self.nexts = []
        
    def add_test(self, leftFact, leftAttr, rel, rightFact, rightAttr):
        self.logger.debug("NotNode.add_test(leftFact = %s, leftAttr = %s, rel = %s, rightFact = %s, rightAttr = %s)" % (leftFact, leftAttr, rel, rightFact, rightAttr))
        
        self.tests.append((leftFact, leftAttr, rel, rightFact, rightAttr))
        
    def left_activate(self, tag, timestamp, token):
        self.logger.debug("NotNode.left_activate(tag = %s, timestamp = %s, token = %s)" % (tag, timestamp, token))
        
        for fact in self.right.tokens:
            failed = False
            for leftFact, leftAttr, rel, rightFact, rightAttr in self.tests:
                leftValue = getattr(token[leftFact], leftAttr)
                rightValue = getattr(fact, rightAttr)
                
                if not self.testfunc[rel](leftValue, rightValue):
                    failed = True
                    break
                    
            if not failed:
                if tag == '+':
                    try:
                        self.count[token.id] += 1
                        
                    except KeyError:
                        self.count[token.id] = 1
                        
                elif tag == '-':
                    del self.count[token.id]
                    
            if self.count[token.id] == 0:
                t = copy.copy(token)
                t[rightFact] = fact
                
                for next in self.nexts:
                    next.left_activate(tag, timestamp, t)
                    
    def right_activate(self, tag, timestamp, fact):
        self.logger.debug("NotNode.right_activate(tag = %s, timestamp = %s, fact = %s)" % (tag, timestamp, fact))
        
        for token in self.left.tokens:
            failed = False
            for leftFact, leftAttr, rel, rightFact, rightAttr in self.tests:
                leftValue = getattr(token[leftFact], leftAttr)
                rightValue = getattr(fact, rightAttr)
                
                if not self.testfunc[rel](leftValue, rightValue):
                    failed = True
                    break
                    
            if not failed:
                if tag == '+':
                    try:
                        self.count[token.id] += 1
                        
                    except KeyError:
                        self.count[token.id] = 1
                        
                    if self.count[token.id] == 1:
                        t = copy.copy(token)
                        t[rightFact] = fact
                        
                        for next in self.nexts:
                            next.left_activate('-', timestamp, t)
                            
                elif tag == '-':
                    self.count[token.id] -= 1
                    
                    if self.count[token.id] == 0:
                        t = copy.copy(token)
                        t[rightFact] = fact
                        
                        for next in self.nexts:
                            next.left_activate('+', timestamp, t)
                            
                            
class Memory(Node):
    def __init__(self):
        self.logger = logging.getLogger('PyRete')
        self.logger.debug("Memory.__init__()")
        
        self.nexts = []
        self.tokens = []
        
    def left_activate(self, tag, timestamp, token):
        self.logger.debug("Memory.left_activate(tag = %s, timestamp = %s, token = %s)" % (tag, timestamp, token))
        
        if tag == '+':
            self.tokens.append(token)
            
        elif tag == '-':
            try:
                self.tokens.remove(token)
                
            except ValueError:
                pass
                
        for next in self.nexts:
            next.left_activate(tag, timestamp, token)
            
    def right_activate(self, tag, timestamp, fact):
        self.logger.debug("Memory.right_activate(tag = %s, timestamp = %s, fact = %s)" % (tag, timestamp, fact))
        
        if tag == '+':
            self.tokens.append(fact)
            
        elif tag == '-':
            self.tokens.remove(fact)
            
        for next in self.nexts:
            next.right_activate(tag, timestamp, fact)
            
            
class ProductionNode(Node):
    def __init__(self, rule):
        self.logger = logging.getLogger('PyRete')
        self.logger.debug("ProductionNode.__init__(rule = %s)" % (rule))
        
        self.rule = rule
        self.doc = None
        
        self.salience = 0
        self.autofocus = False
        
        self.env = None
        self.rhs = None
        
        self.bindings = {}
        
        self.tokens = []
        self.executed_tokens = []
        
    def left_activate(self, tag, timestamp, token):
        self.logger.debug("ProductionNode.left_activate(tag = %s, timestamp = %s, token = %s)" % (tag, timestamp, token))
        
        if tag == '+':
            self.tokens.append(token)
            
        elif tag == '-':
            try:
                self.tokens.remove(token)
                
            except ValueError:
                pass
                
            
    def right_activate(self, tag, timestamp, fact):
        self.logger.debug("ProductionNode.right_activate(tag = %s, timestamp = %s, fact = %s)" % (tag, timestamp, fact))
        
        if tag == '+':
            self.tokens.append(fact)
            
        elif tag == '-':
            self.tokens.remove(fact)
            
    def remove_token(self, token):
        try:
            self.tokens.remove(token)
            self.executed_tokens.append(token)
            
        except ValueError:
            pass
            
            
class Rete(object):
    def __init__(self, rule_engine):
        self.logger = logging.getLogger('PyRete')
        self.rule_engine = rule_engine
        
        self.root_node = {}         # type -> [node1, ..., nodeN,]
        self.production_nodes = {}  # ruleName -> ProductionNode
        
    def timestamp(self):
        ts = 0
        while True:
            yield ts
            ts += 1
            
    def remove_activation(self, activation):
        """
        """
        
        rule, salience, token, rhs, locals = activation
        self.production_nodes[rule].remove_token(token)
        
    def add_object(self, obj):
        """
        Add an object to the Rete Network.
        """
        
        ts = self.timestamp()
        self.logger.debug("Rete.add_object(obj = %s)" % (obj))
        try:
            for obj_type_test in self.root_node[obj.__class__.__name__]:
                obj_type_test.right_activate('+', ts, obj)
                
        except KeyError:
            pass
            
    def del_object(self, obj):
        """
        Delete an object from the Rete Network.
        """
        
        ts = self.timestamp()
        self.logger.debug("Rete.del_object(obj = %s)" % (obj))
        try:
            for obj_type_test in self.root_node[obj.__class__.__name__]:
                obj_type_test.right_activate('-', ts, obj)
                
        except KeyError:
            pass
            
            
    def get_conflict_set(self, rule = None):
        """
        Returns the conflict set.
        
        The conflict set is generated by joining all tokens from all
        production nodes. The conflict set is used to generate the agenda.
        """
        
        conflict_set = []
        if rule is None:
            for _rule in self.production_nodes:
                conflict_set.append((_rule,
                                     self.production_nodes[_rule].salience,
                                     self.production_nodes[_rule].tokens,
                                     self.production_nodes[_rule].rhs,
                                     self.production_nodes[_rule].bindings),)
                                     
        else:
            conflict_set.extend((rule,
                                 self.production_nodes[rule].salience,
                                 self.production_nodes[rule].tokens,
                                 self.production_nodes[rule].rhs,
                                 self.production_nodes[_rule].bindings))
                                 
        return conflict_set
        
        
    def _connect_nodes(self, current, next, _from = "left"):
        self.logger.debug("Rete._connect_nodes(current = %s, next = %s, _from = %s)" % (current, next, _from))
        
        try:
            current.nexts.append(next)
            if isinstance(next, JoinNode) or \
               isinstance(next, NotNode) or \
               isinstance(next, DummyJoinNode):
                setattr(next, _from, current)
                
            if isinstance(next, LeftInputAdapter):
                setattr(next, _from, current)
                next._tokens = current.tokens
                
        except AttributeError:
            pass
            
    def _add_to_root(self, node, type):
        try:
            self.root_node[type].append(node)
        except:
            self.root_node[type] = [node,]
            
            
    def _add_production(self, rule, node):
        self.production_nodes[rule] = node
        
        